import torch
import torch.nn as nn
from torchinfo import summary
from torch.autograd import Variable
import torch.nn.functional as F
import numpy as np


class MeshClsNet(nn.Module):
    def __init__(self, num_classes=5, num_channels=18):
        super(MeshClsNet, self).__init__()
        self.num_classes = num_classes
        self.num_channels = num_channels

        self.fc1 = nn.Linear(self.num_channels, 18)
        self.fc2 = nn.Linear(18, self.num_classes)

    def forward(self, x, m_seg):
        batchsize = x.size()[0]
        n_pts = x.size()[2]

        x = x.transpose(2, 1) #Q: x[6000*18], K: m_seg[6000*15], V: [6000*15]
        x = self.fc1(x)
        x = x.transpose(2, 1) #x[1000*18], [500*18], --> [1*18], [1*18], -->[15*18]
        x = torch.bmm(x, m_seg)
        x = x.transpose(2, 1)
        x = self.fc2(x)

        x = torch.nn.Softmax(dim=-1)(x.view(-1, self.num_classes))
        x = x.view(batchsize, 15, self.num_classes)
        x = x.transpose(2, 1)
        return x